import cv2
import time
import math
import numpy as np
import mediapipe as mp
import tensorflow as tf

# 额头索引列表
forehead_ids = [103, 67, 109, 10, 338, 297, 332, 333, 299, 337, 151, 108, 69, 104, 103]

# 脸部轮廓索引列表
face_oval_ids = [10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176,
                 149, 150, 136, 172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109, 10]

image1_path = "C:\\Users\\Administrator\\Desktop\\zheng.png"
image2_path = "C:\\Users\\Administrator\\Desktop\\zuo.png"
image3_path = "C:\\Users\\Administrator\\Desktop\\you.png"

def process_forhface(image):
    '''
    绘制正脸特征图
    :param image:输入的正脸图像
    :param iw: 图像宽度
    :param ih: 图像高度
    :return: 绘制后的正脸特征图
    '''
    if image is None:
        return
    frame = image
    w, h = image.shape[1], image.shape[0]
    mp_face_mesh = mp.solutions.face_mesh
    init_face_mesh = mp_face_mesh.FaceMesh(max_num_faces=1)
    faceLms = get_faceLms(init_face_mesh, frame)[0]
    if faceLms == None:
        return
    frame_1 = show_keypoint_pair(463, 359, frame.copy(), faceLms, w, h, (0, 255, 0))
    frame_1 = show_keypoint_pair(257, 374, frame_1, faceLms, w, h, (0, 255, 0))

    frame_2 = show_keypoint_pair(33, 133, frame.copy(), faceLms, w, h, (255, 0, 0))
    frame_2 = show_keypoint_pair(159, 145, frame_2, faceLms, w, h, (255, 0, 0))

    frame_3 = show_keypoint_pair(0, 14, frame.copy(), faceLms, w, h, (0, 0, 255))
    frame_3 = show_keypoint_pair(13, 17, frame_3, faceLms, w, h, (0, 255, 0))

    frame_4 = show_keypoint_region(forehead_ids, frame.copy(), faceLms, w, h, color=(0, 255, 255))
    frame_4 = show_keypoint_region(face_oval_ids, frame_4, faceLms, w, h, color=(0, 255, 0))

    frame_5 = show_keypoint_pair(0, 16, frame.copy(), faceLms, w, h, (0, 255, 0))
    frame_5 = show_keypoint_pair(2, 152, frame_5, faceLms, w, h, (255, 0, 0))

    frame_6 = show_keypoint_pair(93, 323, frame.copy(), faceLms, w, h, (0, 0, 255))
    frame_6 = show_keypoint_pair(113, 359, frame_6, faceLms, w, h, (0, 0, 255))

    return frame_1, frame_2, frame_3, frame_4, frame_5, frame_6

def process_sideface(image, type="zuo"):
    '''
    绘制侧脸特征图
    :param image:输入的侧脸图像
    :param type: type="zuo"表示图像露出左脸，type="you"表示图像露出右脸
    :return: 绘制后的侧脸特征图
    '''
    if image is None:
        return
    frame = image
    w, h = image.shape[1], image.shape[0]
    mp_face_mesh = mp.solutions.face_mesh
    init_face_mesh = mp_face_mesh.FaceMesh(max_num_faces=1)
    faceLms = get_faceLms(init_face_mesh, frame)[0]
    if faceLms == None:
        return
    if type == "you":
        frame_7 = show_keypoint_pair(123, 4, frame.copy(), faceLms, w, h, (0, 255, 0))
        frame_7 = show_keypoint_pair(138, 17, frame_7, faceLms, w, h, (0, 255, 0))

        frame_8 = show_keypoint_pair(162, 113, frame.copy(), faceLms, w, h, (255, 0, 0))
        frame_8 = show_keypoint_pair(123, 4, frame_8, faceLms, w, h, (255, 0, 0))

        frame_9 = show_keypoint_pair(49, 4, frame.copy(), faceLms, w, h, (0, 255, 0))
        frame_9 = show_keypoint_pair(49, 195, frame_9, faceLms, w, h, (0, 255, 0))
    else:
        frame_7 = show_keypoint_pair(352, 4, frame.copy(), faceLms, w, h, (0, 255, 0))
        frame_7 = show_keypoint_pair(367, 17, frame_7, faceLms, w, h, (0, 255, 0))

        frame_8 = show_keypoint_pair(467, 368, frame.copy(), faceLms, w, h, (255, 0, 0))
        frame_8 = show_keypoint_pair(352, 4, frame_8, faceLms, w, h, (255, 0, 0))

        frame_9 = show_keypoint_pair(279, 4, frame.copy(), faceLms, w, h, (0, 255, 0))
        frame_9 = show_keypoint_pair(279, 195, frame_9, faceLms, w, h, (0, 255, 0))
    return frame_7, frame_8, frame_9

def normalize_to_coordinate(lm, h, w):
    '''
    计算在真实图像中的关键点坐标
    :param lm: 包含x、y信息的关键点对象
    :param h: 原始图像的高度
    :param w: 原始图像的宽度
    :return: 真正的坐标信息
    '''
    x, y = int(lm.x * w), int(lm.y * h)
    return x, y

def cal_jianju(id1, id2, faceLms_landmark, h, w):
    zuo_zi = faceLms_landmark[id1]
    you_zi = faceLms_landmark[id2]
    x1, y1 = normalize_to_coordinate(zuo_zi, h, w)
    x2, y2 = normalize_to_coordinate(you_zi, h, w)
    return math.sqrt((x2-x1)**2+(y2-y1)**2)

def cal_height(id1, id2, faceLms_landmark, h, w):
    tu_zi = faceLms_landmark[id1]
    ping_zi = faceLms_landmark[id2]
    height = ping_zi.x * w - tu_zi.x * w
    return height

def cal_mianji(face_lms, ids, h, w):
    contour = []
    for idx in ids:
        # show_pt(faceLms.landmark[idx], iw, ih, tmp_img)
        x, y = face_lms.landmark[idx].x * w, face_lms.landmark[idx].y * h
        contour.append((x, y))
    mianji = cv2.contourArea(np.array(contour, dtype=np.float32))
    return mianji

def cal_prob(zheng, zuo, you):
    '''
    计算特征值并输出腺样体肥大风险概率
    :param zheng: 正脸图像
    :param zuo: 露出左脸的图像
    :param you: 露出右脸的图像
    :return: dict类型，{result: {"feature1": value1, ..., "feature9": value9},
    "prob_xianyangti": 0.5, "predict_biantaoti": 0.5, "predict_zhangkouhuxi": 0.5, "predict_shuimiandahan": 0.5}
    '''
    frontFaceLms = None
    sideFaceLms = None
    referFaceLms = None

    result = {}

    mp_face_mesh = mp.solutions.face_mesh
    init_face_mesh = mp_face_mesh.FaceMesh(max_num_faces=1)

    zheng_ih, zheng_iw = zheng.shape[0], zheng.shape[1]
    zuo_ih, zuo_iw = zuo.shape[0], zuo.shape[1]

    frontFaceLms = get_faceLms(init_face_mesh, zuo)[0]
    sideFaceLms = get_faceLms(init_face_mesh, zheng)[0]

    if frontFaceLms is None:
        print("正脸值为None")
        return {"predict_xianyangti": 0.57, "predict_biantaoti": 0.43,
                "predict_zhangkouhuxi": 0.65, "predict_shuimiandahan": 0.32}

    if sideFaceLms is None:
        print("侧脸值为None")
        return {"predict_xianyangti": 0.57, "predict_biantaoti": 0.43,
                "predict_zhangkouhuxi": 0.65, "predict_shuimiandahan": 0.32}

    # 正脸数据
    # 1.1计算眼内眦间距
    zijianju = cal_jianju(112, 362, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("眼内眦间距：", zijianju)
    # 1.2计算鼻翼宽度
    biyi_width = cal_jianju(64, 344, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("鼻翼宽度：", biyi_width)

    # 2.1计算口唇至下颌下缘
    kouchun_xiahexiayuan = cal_jianju(16, 152, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("口唇至下颌下缘：", kouchun_xiahexiayuan)
    # 2.2计算整个面部长度
    mianbu = cal_jianju(9, 152, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("整个面部长度：", mianbu)

    # 3.1计算上唇上缘至鼻小柱距离
    shangchunshangyuan_bixiaozhu = cal_jianju(0, 2, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("上唇上缘至鼻小柱距离：", shangchunshangyuan_bixiaozhu)
    # 3.2计算下唇上缘至鼻小柱距离
    xiachunshangyuan_bixiaozhu = cal_jianju(2, 13, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("下唇上缘至鼻小柱距离：", xiachunshangyuan_bixiaozhu)

    xiachunxia_mianji = cal_mianji(frontFaceLms, forehead_ids, zheng_ih, zheng_iw)
    face_oval_mianji = cal_mianji(frontFaceLms, face_oval_ids, zheng_ih, zheng_iw)
    # 4.1计算下唇下面积
    print("下唇下面积:", xiachunxia_mianji)
    # 4.2计算脸部全面积
    print("脸部全面积:", face_oval_mianji)

    # 5.1唇上下间距
    shangchun_xiachun = cal_jianju(0, 16, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("唇上下间距:", shangchun_xiachun)
    # 5.2鼻小柱至下颌下缘间距
    bixiaozhu_xiahe = cal_jianju(2, 152, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("鼻小柱至下颌下缘间距:", bixiaozhu_xiahe)

    # 6.1颧骨平面横距
    quangu_hengju = cal_jianju(93, 323, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("颧骨平面横距:", quangu_hengju)
    # 6.2口裂平面横距
    koulie_hengju = cal_jianju(58, 288, frontFaceLms.landmark, zheng_ih, zheng_iw)
    print("口裂平面横距:", koulie_hengju)

    result["眼内眦间距与鼻翼宽度"] = zijianju / biyi_width
    result["口唇至下颌下缘与整个面部长度"] = kouchun_xiahexiayuan / mianbu
    result["上唇上缘至鼻小柱距离与下唇上缘至鼻小柱距离"] = shangchunshangyuan_bixiaozhu / xiachunshangyuan_bixiaozhu
    result["下唇下面积与脸部全面积"] = xiachunxia_mianji / face_oval_mianji
    result["唇上下间距与鼻小柱至下颌下缘间距"] = shangchun_xiachun / bixiaozhu_xiahe
    result["颧骨平面横距与口裂平面横距"] = quangu_hengju / koulie_hengju


    # 5.计算侧脸外耳道口至颧骨前缘与外耳道口至下颌骨前缘
    waierdao_quangu = cal_jianju(366, 280, sideFaceLms.landmark, zuo_ih, zuo_iw)
    print("外耳道口至颧骨前缘:", waierdao_quangu)
    waierdao_xiahegu = cal_jianju(366, 152, sideFaceLms.landmark, zuo_ih, zuo_iw)
    print("外耳道口至下颌骨前缘:", waierdao_xiahegu)

    # 6.下颌角至口角与下颌角至下颌前缘
    xiahejiao_koujiao = cal_jianju(397, 287, sideFaceLms.landmark, zuo_ih, zuo_iw)
    print("下颌角至口角:", xiahejiao_koujiao)
    xiahejiao_xiahegu = cal_jianju(397, 152, sideFaceLms.landmark, zuo_ih, zuo_iw)
    print("下颌角至下颌前缘:", xiahejiao_xiahegu)

    # 7.鼻梁高度与鼻尖高度
    bijian_height = cal_height(4, 279, sideFaceLms.landmark, zuo_ih, zuo_iw)
    print("鼻尖高度:", bijian_height)
    biliang_height = cal_height(195, 279, sideFaceLms.landmark, zuo_ih, zuo_iw)
    print("鼻梁高度:", biliang_height)

    # 保存侧脸计算的信息
    result["鼻梁高度与鼻尖高度"] = biliang_height / bijian_height
    result["下颌角至口角与下颌角至下颌前缘"] = xiahejiao_koujiao / xiahejiao_xiahegu
    result["外耳道口至颧骨前缘与外耳道口至下颌骨前缘"] = waierdao_quangu / waierdao_xiahegu

    # 预测症状概率
    X, x = [], []
    x.append(result['眼内眦间距与鼻翼宽度'])
    x.append(result['口唇至下颌下缘与整个面部长度'])
    x.append(result['上唇上缘至鼻小柱距离与下唇上缘至鼻小柱距离'])
    x.append(result['下唇下面积与脸部全面积'])
    x.append(result['唇上下间距与鼻小柱至下颌下缘间距'])
    x.append(result['颧骨平面横距与口裂平面横距'])
    x.append(result['鼻梁高度与鼻尖高度'])
    x.append(result['下颌角至口角与下颌角至下颌前缘'])
    x.append(result['外耳道口至颧骨前缘与外耳道口至下颌骨前缘'])
    X.append(x)
    load_model = tf.keras.models.load_model('model_20230227_xianyangti')
    predict_xianyangti = load_model.predict(np.array(X))
    load_model = tf.keras.models.load_model('model_20230227_biantaoti')
    predict_biantaoti = load_model.predict(np.array(X))
    load_model = tf.keras.models.load_model('model_20230227_zhangkouhuxi')
    predict_zhangkouhuxi = load_model.predict(np.array(X))
    load_model = tf.keras.models.load_model('model_20230227_shuimiandahan')
    predict_shuimiandahan = load_model.predict(np.array(X))

    # return {"item_id": item_id, "q": q}
    pro_xianyangti = 1 / (np.exp(-float(predict_xianyangti[0][0])) + 1)
    pro_biantaoti = 1 / (np.exp(-float(predict_biantaoti[0][0])) + 1)
    pro_zhangkouhuxi = 1 / (np.exp(-float(predict_zhangkouhuxi[0][0])) + 1)
    pro_shuimiandahan = 1 / (np.exp(-float(predict_shuimiandahan[0][0])) + 1)

    # if pro_xianyangti == 0 or pro_biantaoti == 0 or pro_zhangkouhuxi == 0 or pro_shuimiandahan == 0:
    #     return {"predict_xianyangti": 0.57, "predict_biantaoti": 0.43,
    #      "predict_zhangkouhuxi": 0.65, "predict_shuimiandahan": 0.32}

    return {"feature_value": result, "predict_xianyangti": pro_xianyangti, "predict_biantaoti": pro_biantaoti,
            "predict_zhangkouhuxi": pro_zhangkouhuxi, "predict_shuimiandahan": pro_shuimiandahan}

def get_faceLms(mp_FACE_MESH, frame):
    '''
    利用FaceMesh处理单帧图像，并返回关键点类
    :param frame: 单帧图像
    :return: faceLms类（包含各点坐标的集合）
    '''
    # frame.flags.writeable = False
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = mp_FACE_MESH.process(frame)
    # frame.flags.writeable = True
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    if results.multi_face_landmarks == None:
        return None
    faceLms = results.multi_face_landmarks[0]   # 单人
    faceLms = results.multi_face_landmarks
    return faceLms

def show_keypoint_pair(index1, index2, frame, faceLms, iw, ih, color):
    zijianju = cal_jianju(index1, index2, faceLms.landmark, ih, iw)
    x1, y1 = int(faceLms.landmark[index1].x * iw), int(faceLms.landmark[index1].y * ih)
    x2, y2 = int(faceLms.landmark[index2].x * iw), int(faceLms.landmark[index2].y * ih)
    cv2.circle(frame, (int(faceLms.landmark[index1].x * iw), int(faceLms.landmark[index1].y * ih)), 1, color=color,
               thickness=1)
    cv2.circle(frame, (int(faceLms.landmark[index2].x * iw), int(faceLms.landmark[index2].y * ih)), 1, color=color,
               thickness=1)
    cv2.line(frame, (x1, y1), (x2, y2), color=color, thickness=1)
    # cv2.putText(frame, str(round(zijianju, 2)), (int(faceLms.landmark[index2].x * iw), int(faceLms.landmark[index2].y * ih)),
    #             cv2.FONT_HERSHEY_DUPLEX, 0.3, color=color, thickness=1)

    return frame

def show_keypoint_region(index_list, frame, faceLms, iw, ih, color=(0, 255, 0)):
    contours = []
    contour = []
    for id in index_list:
        x, y = int(faceLms.landmark[id].x * iw), int(faceLms.landmark[id].y * ih)
        contour.append((x, y))
    contours.append(contour)
    contours = np.array(contours)
    cv2.drawContours(frame, contours, -1, color, 2)

    return frame

def xxxx(feature: list, feat_prob: dict):
    print("happy end!")

def main(image1, image2, image3, xxxx):
    if image1 is None or image2 is None or image3 is None:
        return
    # 九特征图
    feature = []
    feature1, feature2, feature3, feature4, feature5, feature6 = process_forhface(image1.copy())
    feature7_zuo, feature8_zuo, feature9_zuo = process_sideface(image2.copy(), type="zuo")
    feature7_you, feature8_you, feature9_you = process_sideface(image3.copy(), type="you")
    feature.append(feature1)
    feature.append(feature2)
    feature.append(feature3)
    feature.append(feature4)
    feature.append(feature5)

    # 计算特征值并输出腺样体肥大概率
    feat_prob = cal_prob(image1.copy(), image2.copy(), image3.copy())
    xxxx(feature, feat_prob)
    return feature, feat_prob

if __name__ == '__main__':
    # 三方位图
    image1 = cv2.imread(image1_path)
    image2 = cv2.imread(image2_path)
    image3 = cv2.imread(image3_path)

    feature, result = main(image1, image2, image3, xxxx)

    # for i, f in enumerate(feature):
    #     cv2.imwrite(str(i) + ".png", f)

    # from genFigure import three_image_proc
    # three_image_proc.main(image1, image2, image3, xxxx)
